import cv2 as cv
import numpy as np
import concurrent.futures

try:
    from rapidocr_onnxruntime import RapidOCR
except:
    from rapidocr_paddle import RapidOCR
from service.QrCodeHandlerService import *
from config.ServerApp import app
from service.TextClassifyHadler import ClassifyHandlerBase
import time
from utils.ImgUtils import *
from utils.Singleton import Singleton
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True


@Singleton
class OcrModel:
    def __init__(self):
        if hasattr(self, 'initialized'):
            return
        self.initialized = True
        from config.Common import get_ocr_device
        import time

        ocr_device = get_ocr_device()
        print(f"[OCR模型] 开始初始化，设备: {ocr_device}")

        load_start = time.time()
        if ocr_device == 'cuda' or ocr_device == 'mlu':
            self.model = RapidOCR(use_cuda=True, det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
            print(f"[OCR模型] 使用GPU加速初始化")
        else:
            self.model = RapidOCR()
            print(f"[OCR模型] 使用CPU初始化")

        load_time = time.time() - load_start
        print(f"[OCR模型] 初始化完成，耗时: {load_time:.3f}秒，设备: {ocr_device}")

    def __call__(self, image_path):
        return self.model(image_path)


__rapid_ocr = OcrModel()


def __get_all_classify_handler():
    '''
    获取所有分类处理handler
    '''
    subclasses = ClassifyHandlerBase.__subclasses__()
    all_subclasses = list(subclasses)
    instances = [subclass(app) for subclass in all_subclasses]
    return instances


def classfy_by_ocr(local_image_path: str):
    '''
    通过ocr分类
    卡证	身份证、银行卡、驾驶证、户口本、护照、营业执照
    文档	支付截图、转账截图、短信验证码截图、聊天截图、快递单、股票截图、体育彩票截图、航班截图、行程订票截图


    票据、红头公章文件、名片、六合彩截图、

    返回数组格式[{},{}]
    '''
    # ocr结果提取
    start_time = time.time()
    ocr_result = ocr_extract(local_image_path)
    end_time = time.time()
    ocr_cost = end_time - start_time
    app.ctx.logger.info(f'===============》OCR耗时：{ocr_cost}')
    print(f'===============》OCR耗时：{ocr_cost}')
    print(f"{local_image_path},OCR文本：{ocr_result}")
    # ocr 分类
    c_start_time = time.time()
    ocr_classfy_result = ocr_classfy(ocr_result)
    c_end_time = time.time()
    ocr_classfy_cost = c_end_time - c_start_time
    print(f'==================>ocr分类耗时：{ocr_classfy_cost}')
    app.ctx.logger.info(f'==================>ocr分类耗时：{ocr_classfy_cost}')
    return ocr_classfy_result, ocr_result


def ocr_classfy(str_input):
    '''
    ocr分类

    卡证:身份证、银行卡、户口簿、驾驶证、中国护照、营业执照、通行证、
    文档：票据、支付截图、转账截图、短信验证码截图、红头公章文件、聊天截图、快递单、名片、股票截图、体育彩票截图、六合彩截图、航班截图、行程订票截图

    返回：一张图片命中得所有分类
    [
        {
            "classifyName": "卡证",
            "subClassifyName": "身份证1",
            "confidence": 1.0
        },
        {
            "classifyName": "卡证",
            "subClassifyName": "身份证2",
            "confidence": 1.0
        }
    ]
    '''

    ocr_classfy_result = []
    if (str_input is None or len(str_input) == 0):
        return ocr_classfy_result
    # 去掉空格 【6283 5065 3600 0129 - 6283506536000129】
    no_blank_str = str_input.replace(" ", "")
    no_blank_str = str_input.replace("\n", "")
    handlers = __get_all_classify_handler()
    # with concurrent.futures.ThreadPoolExecutor(max_workers=len(handlers)) as executor:
    #     futures = []
    for handler in handlers:
        result = handler.verify(no_blank_str)
        # 等待处理完毕
        # for future in concurrent.futures.as_completed(futures):
        #     result = future.result()
        if result['confidence'] == 1:
            ocr_classfy_result.append(result)
    return ocr_classfy_result

def retinex_scales_distribution(size):
    return np.array([2, 8, 16, 32], dtype=np.float32)

def single_scale_retinex(img, sigma):
    retinex = np.log10(img + 1.0) - np.log10(cv.GaussianBlur(img, (0, 0), sigma) + 1.0)
    return retinex

def multi_scale_retinex(img, scales):
    retinex = np.zeros_like(img, dtype=np.float32)
    for sigma in scales:
        retinex += single_scale_retinex(img, sigma)
    retinex = retinex / len(scales)
    return retinex

def MSRCR(img, scales, gain=1.0, offset=0.0):
    img = np.float32(img) + 1.0
    retinex = multi_scale_retinex(img, scales)
    for i in range(img.shape[2]):
        img[:, :, i] = gain * (retinex[:, :, i] - np.min(retinex[:, :, i])) / (
                    np.max(retinex[:, :, i]) - np.min(retinex[:, :, i])) * 255.0 + offset
    img = np.uint8(np.clip(img, 0, 255))
    return img

    img = Image.open(image_path)
    img = img.resize((640, 480))  # 调整图像大小
    img = img.convert("L")  # 灰度化
    img = img.point(lambda p: 0 if p < 128 else 255)  # 二值化
    return img


def ocr_extract(image_path):
    text = ""
    # 读取图片文件，并将其转换为OpenCV图像格式
    try:
        import time
        ocr_start = time.time()

        print(f"[OCR处理] 开始处理图片: {image_path}")

        # 判断是否是本地文件
        # if not os.path.isfile(image_path):
        #     image_path = download_file(image_path)
        # 去除阴影
        img = cv.imread(image_path)
        if img is None:
            print(f"[OCR处理] 图片读取失败: {image_path}")
            return text

        # rapidocr
        result = __rapid_ocr(img)
        result = [d for d in result if d is not None]

        ocr_time = time.time() - ocr_start

        if result is None or len(result) == 0:
            print(f"[OCR处理] 未识别到文本，耗时: {ocr_time:.3f}秒")
            return text

        result_content = result[0]
        text_lines = 0
        for line in enumerate(result_content):
            if line[1][1] != "":
                text += line[1][1] + "\n"
                text_lines += 1

        print(f"[OCR处理] 识别完成，耗时: {ocr_time:.3f}秒，识别行数: {text_lines}")
        return text
    except Exception as e:
        ocr_time = time.time() - ocr_start if 'ocr_start' in locals() else 0
        print(f"[OCR处理] 处理异常，耗时: {ocr_time:.3f}秒，错误: {str(e)}")
        app.ctx.logger.error(f"图像文本解析分类失败path：{image_path},error:{str(e)}")
        return text





